package com.snailcatmall.www.window;

import com.snailcatmall.www.App;
import com.snailcatmall.www.callback.CallBackData;
import com.snailcatmall.www.tools.AllTools;
import javafx.application.Platform;
import javafx.event.EventHandler;
import javafx.scene.control.Button;
import javafx.scene.input.MouseEvent;

import java.io.File;
import java.io.FileReader;
import java.io.IOException;

public class UpdateModelClick {
    public static Button button = new Button();
    public static File file1;
    public static File file2;
    public static String modelType;
    public static String[] modelSig = {"string_to_param", "model.diffusion_model.", "cond_stage_model.transformer.", "lora_te_text_model_encoder", "encoder.down.0.block", "linear.0.weight", "linear1.weight"};
    public static String[] paths = {"\\embeddings\\","\\models\\Stable-diffusion\\","\\models\\Lora\\","\\models\\VAE\\", "\\models\\hypernetworks\\"};
//public static String[] paths = {"\\models\\VAE\\"};
    static {
        init();
    }

    public static void init() {
        button.setLayoutX(160);
        button.setLayoutY(170);
        button.setPrefWidth(150);
        button.setPrefHeight(60);
        button.setText("Update Model");
        button.setOnMouseClicked(new EventHandler<MouseEvent>() {
            @Override
            public void handle(MouseEvent event) {
                ListViewFiles.list.clear();
                Runnable runnable = new Runnable() {
                    @Override
                    public void run() {
                        for (String p:paths){
                            System.out.println(p);
                            button.setDisable(true);
                            boolean flag =  getFiles(p);
                            if (flag){
                                Platform.runLater(new Runnable() {
                                    @Override
                                    public void run() {
                                        App.stage.setTitle("处理 " + p);
                                    }
                                });
                            }
                        }
                        Platform.runLater(new Runnable() {
                            @Override
                            public void run() {
                                App.stage.setTitle("处理完毕！ ");
                                button.setDisable(false);
                            }
                        });

                    }
                };
                Thread thread = new Thread(runnable);
                thread.start();

            }
        });
    }
    public static boolean getFiles( String ps){
        file1 = new File(TextFieldPath.textField.getText() +ps);
        AllTools.recursiveTraversal(file1, new CallBackData() {
            @Override
            public void dataPath(String data) {
                super.dataPath(data);
                file2 = new File(data);
                AllTools.readTextFile(file2, new CallBackData() {
                    @Override
                    public void dataText(String data,FileReader reader) throws IOException {
                        reader.close();
                        super.dataText(data);
                        if (file2.length() < (1024 * 1024) && data.indexOf("string_to_param") != -1) {
                            modelType = "Embedding";
                        } else {
                            for(int i = 0;i< modelSig.length;i++){
                                if (data.indexOf(modelSig[i])!=-1) {
                                    modelType = modelSig[i];
                                    break;
                                }
                            }
                        }
                        if (modelType != null) {
                            Platform.runLater(new Runnable() {
                                @Override
                                public void run() {
                                    ListViewFiles.list.add(file2.getPath());
                                }
                            });
                            File path = file2;
                            boolean flages;
                            System.out.println(data.indexOf("lora_te_text_model_encoder"));
                            if (modelType.equals("string_to_param")) {
                                flages = AllTools.cutFile(path,new File(("\\embeddings\\")));
                            }
                            if (modelType.equals("model.diffusion_model.")) {
                                flages = AllTools.cutFile( path,new File(TextFieldPath.textField.getText() + "\\models\\Stable-diffusion\\"));
                            }
                            if (modelType.equals("cond_stage_model.transformer.")) {
                                flages = AllTools.cutFile( path,new File(TextFieldPath.textField.getText() + "\\models\\Stable-diffusion\\"));
                            }
                            if (modelType.equals("lora_te_text_model_encoder")) {
                                flages = AllTools.cutFile( path, new File(TextFieldPath.textField.getText()  + "\\models\\Lora\\"));
                            }
                            if (modelType.equals("encoder.down.0.block")) {
                                flages = AllTools.cutFile( path, new File(TextFieldPath.textField.getText()  + "\\models\\VAE\\"));
                            }
                            if (modelType.equals("linear.0.weight")) {
                                flages = AllTools.cutFile( path, new File(TextFieldPath.textField.getText()  + "\\models\\hypernetworks\\"));
                            }
                            if (modelType.equals("linear1.weight")) {
                                flages = AllTools.cutFile( path, new File(TextFieldPath.textField.getText()  + "\\models\\hypernetworks\\"));
                            }
                        }

                    }
                });
            }
        });
        return true;
    }

}
